import time

from src.matting.effnetv2 import extract_mask
# from src.video_composition.cv2_composition import video_human_segmentation
from src.video_composition.ffmpeg_composition import video_human_segmentation


def main():
    # video_input_path = 'https://modelscope.oss-cn-beijing.aliyuncs.com/test/videos/video_matting_test.mp4'
    # matting_out_path = 'data/matting/matting_out.mp4'
    # video_out_path = 'data/matting/test_3.mp4'

    video_input_path = "data/ori_video/c.mp4" # data/ori_video/飞书20241011-201502.mp4
    matting_out_path = 'data/matting/c_matting.mp4'

    # 提取视频中的人像
    # result = extract_mask(video_input_path, matting_out_path)
    # result = read_video(matting_out_path, gray=True)

    # 人物和背景视频合成
    # video_out_path = 'data/output/a.mp4'
    # video_human_segmentation(video_input_path, video_out_path, result, back_pic="data/background/mask_backaround.png")
    blur_radius = 0
    video_out_path = f'data/output/c_smooth_blur{blur_radius}.mp4'
    video_human_segmentation(video_input_path, video_out_path, matting_out_path, back_pic="data/background/mask_backaround.png",
                             back_video_path="data/background/mask_backaround.mp4", blur_radius=blur_radius)


def main1():

    video_input_path = "data/ori_video/飞书20241014-172537.mp4" # data/ori_video/飞书20241011-201502.mp4
    matting_out_path = 'data/matting/飞书20241014-172537_matting.mp4'

    start = time.time()
    # 提取视频中的人像
    # result = extract_mask(video_input_path, matting_out_path)
    # result = read_video(matting_out_path, gray=True)
    end = time.time()
    print(f"extract_mask spend time {end -start}s")
    # 人物和背景视频合成
    # video_out_path = 'data/output/a.mp4'
    # video_human_segmentation(video_input_path, video_out_path, result, back_pic="data/background/mask_backaround.png")
    blur_radius = 0
    video_out_path = f'data/output/飞书20241014-172537_1_smooth_blur{blur_radius}.mp4'
    video_human_segmentation(video_input_path, video_out_path, matting_out_path, back_pic="data/background/mask_backaround.png",
                             back_video_path="data/background/mask_backaround.mp4", blur_radius=blur_radius)

if __name__ == '__main__':
    # main()
    main1()